// Copyright 2022 jeff.li. and/or its affiliates.

#pragma once

#include <memory>
#include <random>
#include <string>
#include <vector>

#include <tbir/pipeline/op_kernel.h>
#include <tbir/runtime/bytes_hash.h>
#include <tbir/runtime/logging.h>
#include <tbir/runtime/runtime_value.h>

namespace tbir::runtime {

    class Graph;

    class OpKernel;

    struct Node;
    struct NodeEntry;

    class Symbol;

    using NodeEntryPtr = std::shared_ptr<NodeEntry>;
    using NodePtr = std::shared_ptr<Node>;

    struct NodeEntry {
        NodePtr node = nullptr;
        uint32_t index = 0;
        String key;
        bool exported = false;
        runtime::RTValue data;

        NodeEntry() = default;

        NodeEntry(NodePtr node, uint32_t index, String key) {
            init(std::move(node), index, std::move(key));
        }

        void init(NodePtr n, uint32_t i, String k) {
            this->node = std::move(n);
            this->index = i;
            this->key = std::move(k);
        }

        NodeEntry(NodePtr node, uint32_t index, uint64_t node_sig) {
            init(std::move(node), index, node_sig);
        }

        void init(NodePtr n, uint32_t i, uint64_t n_sig) {
            this->node = std::move(n);
            this->index = i;
            char buf[256];
            int len = snprintf(buf, 256, "%llu_%u", n_sig, i);
            this->key.assign(buf, len);
        }

        bool operator==(const NodeEntry &o) const {
            return node == o.node && index == o.index;
        }

        String Name();
    };

// Symbol is only used for save NodeEntryPtr
    class Symbol {
    public:
        Symbol() : entry_(nullptr), all_entries_() {
        }

        explicit Symbol(const std::vector<NodeEntryPtr> &all_outputs, int64_t output_idx)
                : entry_(all_outputs[output_idx]), all_entries_(all_outputs) {
        }

        ~Symbol() = default;

    public:
        bool operator==(const Symbol &o) const {
            return entry_ == o.entry_;
        }

        bool operator!=(const Symbol &o) const {
            return entry_ != o.entry_;
        }

        const NodeEntryPtr &GetEntry() const {
            return entry_;
        }

        const std::vector<NodeEntryPtr> &GetAllEntries() const {
            return all_entries_;
        }

    private:
        NodeEntryPtr entry_ = nullptr;
        std::vector<NodeEntryPtr> all_entries_;

        friend class Graph;

        friend class TXSession;
    };

    struct NodeOutput {
        NodeEntry *source = nullptr;
        std::weak_ptr<NodeEntry> weak_ref;
    };

    inline std::vector<NodeOutput> MakeNodeOutputs(const std::vector<NodeEntryPtr> &entrys) {
        std::vector<NodeOutput> nos;
        nos.reserve(entrys.size());
        for (auto &ne : entrys) {
            NodeOutput o;
            o.source = ne.get();
            o.weak_ref = ne;
            nos.push_back(std::move(o));
        }
        return std::move(nos);
    }

    struct Node {
        OpKernelPtr op = nullptr;
        std::vector<NodeEntryPtr> inputs;
        std::vector<NodeOutput> outputs;
        std::vector<NodeEntryPtr> holder;
        String name;

        Node() {
            op = nullptr;
            name.reserve(128);
        }

        bool IsVariable() const;

        Dict ToDict() const;

        static NodePtr FromDict(const Dict &config, Graph *g);

        static NodePtr Create() {
            return std::make_shared<Node>();
        }
    };

}  // namespace tbir::runtime
